import numpy as np
from ray.rllib.evaluation.sample_batch_builder import SampleBatchBuilder
from ray.rllib.models.preprocessors import get_preprocessor
from ray.rllib.offline.json_writer import JsonWriter
from tqdm import tqdm

from astar_actor import AstarActor
from uav_2d import Uav2DEnv
from uav_2d.wrappers.raster_wrapper import RasterWrapper

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--max_eps", type=int, default=100)
    args = parser.parse_args()

    batch_builder = SampleBatchBuilder()  # or MultiAgentSampleBatchBuilder
    writer = JsonWriter("astar-out", compress_columns=[])

    env = RasterWrapper(Uav2DEnv())
    actor = AstarActor(env, step_limit=2000)
    prep = get_preprocessor(env.observation_space)(env.observation_space)
    print("The preprocessor is", prep)

    try:
        for eps_id in tqdm(range(args.max_eps)):
            obs, info = env.reset()
            actor.reset()
            prev_action = np.zeros_like(env.action_space.sample())
            prev_reward = 0
            terminated = truncated = False
            t = 0
            while not terminated and not truncated:
                action = actor.get_action()
                new_obs, rew, terminated, truncated, info = env.step(action)
                batch_builder.add_values(
                    t=t,
                    eps_id=eps_id,
                    agent_index=0,
                    obs=prep.transform(obs),
                    actions=action,
                    action_prob=1.0,
                    action_logp=0.0,
                    rewards=rew,
                    prev_actions=prev_action,
                    prev_rewards=prev_reward,
                    terminateds=terminated,
                    truncateds=truncated,
                    infos={},
                    new_obs=prep.transform(new_obs),
                )
                obs = new_obs
                prev_action = action
                prev_reward = rew
                t += 1
                writer.write(batch_builder.build_and_reset())
    finally:
        env.close()
